from utils import paddle_aux
import paddle
from utils.activations import h_swish


class CoordAtt(paddle.nn.Layer):

    def __init__(self, inp, oup, reduction=32):
        super(CoordAtt, self).__init__()
        self.pool_h = paddle.nn.AdaptiveAvgPool2D(output_size=(None, 1))
        self.pool_w = paddle.nn.AdaptiveAvgPool2D(output_size=(1, None))
        mip = max(8, inp // reduction)
        self.conv1 = paddle.nn.Conv2D(in_channels=inp, out_channels=mip,
            kernel_size=1, stride=1, padding=0)
        self.bn1 = paddle.nn.BatchNorm2D(num_features=mip)
        self.act = h_swish()
        self.conv_h = paddle.nn.Conv2D(in_channels=mip, out_channels=oup,
            kernel_size=1, stride=1, padding=0)
        self.conv_w = paddle.nn.Conv2D(in_channels=mip, out_channels=oup,
            kernel_size=1, stride=1, padding=0)

    def forward(self, x):
        identity = x
        n, c, h, w = tuple(x.shape)
        x_h = self.pool_h(x)
        x_w = self.pool_w(x).transpose(perm=[0, 1, 3, 2])
        y = paddle.concat(x=[x_h, x_w], axis=2)
        y = self.conv1(y)
        y = self.bn1(y)
        y = self.act(y)
        x_h, x_w = paddle_aux.split(x=y, num_or_sections=[h, w], axis=2)
        x_w = x_w.transpose(perm=[0, 1, 3, 2])
        a_h = self.conv_h(x_h).sigmoid()
        a_w = self.conv_w(x_w).sigmoid()
        out = identity * a_w * a_h
        return out
